package com.hjzd.ai.smartpractice.websocket;

import com.google.gson.JsonObject;
import com.google.gson.JsonParser;
import org.java_websocket.client.WebSocketClient;
import org.java_websocket.handshake.ServerHandshake;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import org.springframework.web.socket.BinaryMessage;
import org.springframework.web.socket.TextMessage;
import org.springframework.web.socket.WebSocketSession;

import java.io.IOException;
import java.net.URI;
import java.nio.ByteBuffer;
import java.util.Base64;
import java.util.HashMap;
import java.util.Map;
import java.util.concurrent.atomic.AtomicBoolean;

import static com.hjzd.ai.smartpractice.constant.ServiceConstant.FIRST_TTS_PREFIX;
import static com.hjzd.ai.smartpractice.constant.ServiceConstant.TTS_PREFIX;

public class StepfunTTSClient extends WebSocketClient {

    private static final Logger logger = LoggerFactory.getLogger(StepfunTTSClient.class);
    private final WebSocketSession session;
    private String sessionId; // 由服务端返回的 session_id
    private static final String STEP_API_KEY = "hY3zRFDKcKhf30fCX56iM3gZ8cfLWa5DhYhoJCTEGQsId0VZ4xuRXUoWdJW0UtYn";
    private final AtomicBoolean isPlaying = new AtomicBoolean(false); // 标记 TTS 是否正在播放
    private final AtomicBoolean isInterrupted = new AtomicBoolean(false); // 标记是否被打断

    public StepfunTTSClient(URI serverUri, WebSocketSession session) {
        super(serverUri, createHeaders());
        this.session = session;
    }

    // 打断 TTS 播放
    public void interrupt() {
        if (isPlaying.get()) {
            isInterrupted.set(true);
            //logger.info("TTS 播放被打断");
        }
    }

    // 重置打断状态
    public void resetInterrupt() {
        isInterrupted.set(false);
        logger.info("重置打断状态");
    }

    // 创建 WebSocket 请求头，包含鉴权信息
    private static Map<String, String> createHeaders() {
        Map<String, String> headers = new HashMap<>();
        headers.put("Authorization", STEP_API_KEY);
        return headers;
    }

    @Override
    public void onOpen(ServerHandshake handshake) {
        logger.info("阶跃WebSocket连接已打开");
    }

    @Override
    public void onMessage(String message) {
        logger.info("收到阶跃WebSocket消息");
        handleTTSMessage(message);
    }

    @Override
    public void onMessage(ByteBuffer bytes) {
        try {
            if (!isInterrupted.get()) {
                isPlaying.set(true);
                byte[] audioData = bytes.array();
                session.sendMessage(new BinaryMessage(audioData));
                logger.info("已发送阶跃音频数据到前端");
            } else {
                logger.info("TTS 播放被打断，忽略音频数据");
            }
        } catch (IOException e) {
            logger.error("发送阶跃音频数据失败", e);
        } finally {
            isPlaying.set(false);
        }
    }

    @Override
    public void onClose(int code, String reason, boolean remote) {
        logger.info("WebSocket 连接已关闭，原因: {}", reason);
    }

    @Override
    public void onError(Exception ex) {
        logger.error("WebSocket 错误", ex);
    }

    private void handleTTSMessage(String message) {
        JsonObject json = JsonParser.parseString(message).getAsJsonObject();
        String eventType = json.get("type").getAsString();
        if ("tts.connection.done".equals(eventType)) {
            this.sessionId = json.get("data").getAsJsonObject().get("session_id").getAsString();
            logger.info("TTS 连接已完成，session_id: {}", sessionId);
            sendCreateEvent();
        } else if ("tts.response.error".equals(eventType)) {
            logger.error("TTS 错误: {}", message);
            JsonObject data = json.get("data").getAsJsonObject();
            String errorCode = data.get("code").getAsString();
            String errorMessage = data.get("message").getAsString();
            logger.error("错误代码: {}, 错误信息: {}", errorCode, errorMessage);
        } else if ("tts.response.audio.done".equals(eventType)) {
            logger.info("处理TTS音频数据");
            JsonObject data = json.get("data").getAsJsonObject();
            String audioBase64 = data.get("audio").getAsString();
            byte[] audioData = Base64.getDecoder().decode(audioBase64);

            // 输出 type 和 audio 大小
            logger.info("type: {}", eventType);
            logger.info("audio 大小: {} 字节", audioData.length);

            // 发送音频数据到前端
            try {
                //session.sendMessage(new BinaryMessage(audioData));
                logger.info("已发送阶跃音频数据到前端");
            } catch (Exception e) {
                logger.error("发送阶跃音频数据失败", e);
            }
        } else if ("tts.response.audio.delta".equals(eventType)) {
            logger.info("处理 TTS 音频分块数据");
            JsonObject data = json.get("data").getAsJsonObject();
            String audioBase64 = data.get("audio").getAsString();
            byte[] audioData = Base64.getDecoder().decode(audioBase64);

            // 输出 type 和 audio 大小
            logger.info("type: {}", eventType);
            logger.info("audio 大小: {} 字节", audioData.length);

            // 发送音频分块数据到前端
            try {
                session.sendMessage(new BinaryMessage(audioData));
                logger.info("已发送阶跃音频分块数据到前端");
            } catch (IOException e) {
                logger.error("发送阶跃音频分块数据失败", e);
            }
        }
    }


    private void sendCreateEvent() {
        if (sessionId == null) {
            logger.error("session_id 未初始化，无法发送 tts.create 事件");
            return;
        }
        String createEvent = String.format(
                "{\"type\":\"tts.create\",\"data\":{\"session_id\":\"%s\",\"voice_id\":\"voice-tone-FbioybbPhA\",\"response_format\":\"wav\",\"volume_ratio\":1.0,\"speed_ratio\":1.0,\"sample_rate\":16000}}",
                sessionId
        );
        this.send(createEvent);
        logger.info("发送 tts.create 事件: {}", createEvent);
    }

    // 初始化 sessionId
    private boolean initializeSessionId() {
        try {
            // 重新建立 WebSocket 连接
            if (!isOpen()) {
                logger.info("WebSocket 连接已断开，尝试重新连接");
                //reconnect();
            }

            // 等待 tts.connection.done 事件返回 sessionId
            int retryCount = 0;
            while (sessionId == null && retryCount < 3) { // 最多重试 3 次
                logger.info("等待 session_id...");
                Thread.sleep(1000); // 等待 1 秒
                retryCount++;
            }

            if (sessionId == null) {
                logger.error("session_id 获取失败，重试次数已用完");
                return false;
            }

            logger.info("session_id 获取成功: {}", sessionId);
            return true;
        } catch (Exception e) {
            logger.error("初始化 session_id 失败", e);
            return false;
        }
    }

    // 重新建立 WebSocket 连接
    public void reconnect() {
        try {
            if (isOpen()) {
                close();
            }
            connect();
        } catch (Exception e) {
            logger.error("WebSocket 重新连接失败", e);
        }
    }

    // 用于记录当前是否处于括号内的状态
    private boolean inBrackets = false;

    public void sendAudioData(String text) {
        // 过滤掉括号及其中的内容
        if (text == null || text.trim().isEmpty()) {
            logger.warn("检测到空文本，跳过发送");
            return;
        }

        // 检查是否进入或退出括号
        if (text.equals("（") || text.equals("(")) {
            inBrackets = true; // 进入括号
            logger.warn("检测到括号开始，跳过发送");
            return;
        }
        if (text.equals("）") || text.equals(")")) {
            inBrackets = false; // 退出括号
            logger.warn("检测到括号结束，跳过发送");
            return;
        }

        // 如果当前处于括号内，则跳过发送
        if (inBrackets) {
            logger.warn("当前文本在括号内，跳过发送: {}", text);
            return;
        }

        // 检查 sessionId 是否初始化
        if (sessionId == null) {
            logger.warn("session_id 未初始化，尝试重新获取");
            boolean success = initializeSessionId();
            if (!success) {
                logger.error("session_id 获取失败，无法发送音频数据");
                return;
            }
        }

        // 构建音频请求
        String audioRequest = String.format(
                "{\"type\":\"tts.text.delta\",\"data\":{\"session_id\":\"%s\",\"text\":\"%s\"}}",
                sessionId, text
        );

        // 发送音频请求
        this.send(audioRequest);
        logger.info("发送音频请求: {}", audioRequest);

        // 发送文本消息到会话
        try {
            session.sendMessage(new TextMessage(TTS_PREFIX + text));
        } catch (IOException e) {
            logger.error("发送文本消息到会话失败: {}", e.getMessage());
        }
    }

    public void sendFirstAudioData(String text) {
        // 过滤掉括号及其中的内容
        if (text == null || text.trim().isEmpty()) {
            logger.warn("检测到空文本，跳过发送");
            return;
        }

        // 检查是否进入或退出括号
        if (text.equals("（") || text.equals("(")) {
            inBrackets = true; // 进入括号
            logger.warn("检测到括号开始，跳过发送");
            return;
        }
        if (text.equals("）") || text.equals(")")) {
            inBrackets = false; // 退出括号
            logger.warn("检测到括号结束，跳过发送");
            return;
        }

        // 如果当前处于括号内，则跳过发送
        if (inBrackets) {
            logger.warn("当前文本在括号内，跳过发送: {}", text);
            return;
        }

        // 检查 sessionId 是否初始化
        if (sessionId == null) {
            logger.warn("session_id 未初始化，尝试重新获取");
            boolean success = initializeSessionId();
            if (!success) {
                logger.error("session_id 获取失败，无法发送音频数据");
                return;
            }
        }

        // 构建音频请求
        String audioRequest = String.format(
                "{\"type\":\"tts.text.delta\",\"data\":{\"session_id\":\"%s\",\"text\":\"%s\"}}",
                sessionId, text
        );

        // 发送音频请求
        this.send(audioRequest);
        logger.info("发送音频请求: {}", audioRequest);

        // 发送文本消息到会话
        try {
            session.sendMessage(new TextMessage(FIRST_TTS_PREFIX + text));
        } catch (IOException e) {
            logger.error("发送文本消息到会话失败: {}", e.getMessage());
        }
    }

    public void sendTextDone() {
        String doneRequest = String.format(
                "{\"type\":\"tts.text.flush\",\"data\":{\"session_id\":\"%s\"}}",
                sessionId
        );
        this.send(doneRequest);
        // 重置打断状态
        resetInterrupt();
        logger.info("发送tts.text.done请求: {}", doneRequest);
    }
}
